import matplotlib
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import json
import os


def draw_lineChart(data_url="", save_fig=False):
    """
    车辆卸载策略选择:
        0: 本地卸载 (local)
        1: RSU卸载 (rsu)
        2: 随机卸载 (random)
        3: 分布式DDPG卸载 (distributed_DDPG)
        4: 分布式PPO卸载 (distributed_PPO)
        5: 集中式DDPG卸载(联邦学习) (centralized_DDPG)
        6: 集中式PPO卸载(联邦学习) (centralized_PPO)
        else: 未定义, 不可调用
    """
    if data_url == "":
        print("数据目录为空, 无法绘图")
        return
    print(data_url)

    matplotlib.rcParams['font.family'] = 'sans-serif'
    matplotlib.rcParams['font.sans-serif'] = ['SimHei']
    matplotlib.rcParams['axes.unicode_minus'] = False

    # 改为 2x2 子图布局：第一个子图放在左上角
    plt.subplot(2, 2, 1)

    min_y = 100000000000
    max_y = 0

    srs_dict = [["local", "本地卸载", "#bf9000"],
                ["rsu", "RSU卸载", "#7f6000"],
                ["random", "随机卸载", "purple"],
                ["distributed_DDPG", "分布DDPG卸载", "#a9d18e"],
                ["distributed_PPO", "分布PPO卸载", "#8faadc"],
                ["centralized_DDPG", "联邦DDPG卸载", "#548235"],
                ["centralized_PPO", "联邦PPO卸载", "#2f5597"],
                ["distributed_SAC", "分布SAC卸载", "#f4b183"],
                ["centralized_SAC", "集中式SAC卸载", "#c55a11"]]

    for srs in srs_dict:
        if os.path.exists(f'{data_url}/{srs[0]}_return.json'):
            with open(f'{data_url}/{srs[0]}_return.json', 'r') as f:
                y = json.load(f)
            min_y = min(min_y, min(y))
            max_y = max(max_y, max(y))
            x = [i + 1 for i in range(len(y))]
            plt.plot(x, y, color=srs[2], label=srs[1])

    plt.title('训练回报折线图')
    plt.xlabel('训练轮次')
    plt.ylabel('回报')
    plt.xlim(1, len(y) + 1)
    plt.ylim(min_y - 1_000_000, max_y + 1_000_000)
    plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(10))
    # 将第一幅子图的 Y 轴刻度改为自适应，最多 10 个刻度
    plt.gca().yaxis.set_major_locator(ticker.MaxNLocator(10))
    plt.legend(loc='upper right')
    plt.grid()

    # 第二个子图放在右上角（任务卸载成功率）
    plt.subplot(2, 2, 2)

    for srs in srs_dict:
        if os.path.exists(f'{data_url}/{srs[0]}_successRate.json'):
            with open(f'{data_url}/{srs[0]}_successRate.json', 'r') as f:
                y = json.load(f)
            min_y = min(min_y, min(y))
            max_y = max(max_y, max(y))
            x = [i + 1 for i in range(len(y))]
            plt.plot(x, y, color=srs[2], label=srs[1])

    plt.title('任务卸载成功率折线图')
    plt.xlabel('训练轮次')
    plt.ylabel('任务卸载成功率')
    plt.xlim(1, len(y) + 1)
    plt.ylim(0, 100)
    plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(10))
    plt.gca().yaxis.set_major_locator(ticker.MultipleLocator(10))
    plt.legend(loc='upper right')
    plt.grid()

    # 新增：第三个子图（左下角）绘制任务处理平均时延
    plt.subplot(2, 2, 3)

    min_y_latency = 100000000000
    max_y_latency = 0

    for srs in srs_dict:
        if os.path.exists(f'{data_url}/{srs[0]}_taskTotaltime.json'):
            with open(f'{data_url}/{srs[0]}_taskTotaltime.json', 'r') as f:
                y = json.load(f)
            min_y_latency = min(min_y_latency, min(y))
            max_y_latency = max(max_y_latency, max(y))
            x = [i + 1 for i in range(len(y))]
            plt.plot(x, y, color=srs[2], label=srs[1])

    plt.title('任务处理平均时延')
    plt.xlabel('训练轮次')
    plt.ylabel('任务处理平均时延(s)')
    plt.xlim(1, len(y) + 1)
    # 动态设置 y 轴范围，留出 5% 边距
    if max_y_latency > min_y_latency:
        margin = (max_y_latency - min_y_latency) * 0.05
    else:
        margin = 1
    plt.ylim(min_y_latency - margin, max_y_latency + margin)
    plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(10))
    # 让刻度更自适应一些
    plt.gca().yaxis.set_major_locator(ticker.MaxNLocator(10))
    plt.legend(loc='upper right')
    plt.grid()

    # 新增：第四个子图（右下角）绘制任务处理平均能耗
    plt.subplot(2, 2, 4)

    min_y_energy = 100000000000
    max_y_energy = 0

    for srs in srs_dict:
        if os.path.exists(f'{data_url}/{srs[0]}_taskEnergy.json'):
            with open(f'{data_url}/{srs[0]}_taskEnergy.json', 'r') as f:
                y = json.load(f)
            min_y_energy = min(min_y_energy, min(y))
            max_y_energy = max(max_y_energy, max(y))
            x = [i + 1 for i in range(len(y))]
            plt.plot(x, y, color=srs[2], label=srs[1])

    plt.title('任务处理平均能耗')
    plt.xlabel('训练轮次')
    plt.ylabel('任务处理平均能耗（J）')
    plt.xlim(1, len(y) + 1)
    # 动态设置 y 轴范围，留出 5% 边距
    if max_y_energy > min_y_energy:
        margin_e = (max_y_energy - min_y_energy) * 0.05
    else:
        margin_e = 1
    plt.ylim(min_y_energy - margin_e, max_y_energy + margin_e)
    plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(10))
    plt.gca().yaxis.set_major_locator(ticker.MaxNLocator(10))
    plt.legend(loc='upper right')
    plt.grid()

    # 显示图形
    plt.tight_layout()  # 自动调整子图参数以给子图留出足够的空间
    if save_fig:
        plt.savefig(f"{data_url}/line_chart.jpg")
    plt.show()


if __name__ == '__main__':
    draw_lineChart(data_url="../simulation_results/00000_for_draw",
                   save_fig=False)
